"""
工具函数模块，包含数据处理、采样和评估相关的函数
"""

import os
import sys
import copy
import torch
import random
import numpy as np
from collections import defaultdict
from multiprocessing import Process, Queue


def build_index(dataset_name):
    """
    构建用户-物品和物品-用户的索引

    Args:
        dataset_name: 数据集名称

    Returns:
        u2i_index: 用户到物品的映射列表
        i2u_index: 物品到用户的映射列表
    """
    # 获取当前文件的目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    # 构建数据文件的完整路径
    data_file = os.path.join(
        os.path.dirname(current_dir), "data", f"{dataset_name}.txt"
    )

    # 从文本文件加载用户-物品交互矩阵
    ui_mat = np.loadtxt(data_file, dtype=np.int32)

    # 获取最大用户ID和物品ID
    n_users = ui_mat[:, 0].max()
    n_items = ui_mat[:, 1].max()

    # 初始化索引列表
    u2i_index = [[] for _ in range(n_users + 1)]
    i2u_index = [[] for _ in range(n_items + 1)]

    # 构建双向索引
    for ui_pair in ui_mat:
        u2i_index[ui_pair[0]].append(ui_pair[1])
        i2u_index[ui_pair[1]].append(ui_pair[0])

    return u2i_index, i2u_index


def random_neq(l, r, s):
    """
    在指定范围内随机生成一个不在集合s中的整数

    Args:
        l: 随机数下界
        r: 随机数上界
        s: 排除的数字集合

    Returns:
        不在集合s中的随机整数
    """
    t = np.random.randint(l, r)
    while t in s:
        t = np.random.randint(l, r)
    return t


def sample_function(
    user_train, usernum, itemnum, batch_size, maxlen, result_queue, SEED
):
    """
    批量采样函数，用于生成训练数据

    Args:
        user_train: 用户训练数据
        usernum: 用户数量
        itemnum: 物品数量
        batch_size: 批次大小
        maxlen: 序列最大长度
        result_queue: 结果队列
        SEED: 随机种子
    """

    def sample(uid):
        """
        为单个用户采样序列数据

        Args:
            uid: 用户ID

        Returns:
            (uid, seq, pos, neg): 用户ID、序列、正样本、负样本
        """
        while len(user_train[uid]) <= 1:
            uid = np.random.randint(1, usernum + 1)

        seq = np.zeros([maxlen], dtype=np.int32)
        pos = np.zeros([maxlen], dtype=np.int32)
        neg = np.zeros([maxlen], dtype=np.int32)
        nxt = user_train[uid][-1]
        idx = maxlen - 1

        ts = set(user_train[uid])
        for i in reversed(user_train[uid][:-1]):
            seq[idx] = i
            pos[idx] = nxt
            if nxt != 0:
                neg[idx] = random_neq(1, itemnum + 1, ts)
            nxt = i
            idx -= 1
            if idx == -1:
                break

        return (uid, seq, pos, neg)

    np.random.seed(SEED)
    uids = np.arange(1, usernum + 1, dtype=np.int32)
    counter = 0
    while True:
        if counter % usernum == 0:
            np.random.shuffle(uids)
        one_batch = []
        for i in range(batch_size):
            one_batch.append(sample(uids[counter % usernum]))
            counter += 1
        result_queue.put(zip(*one_batch))


class WarpSampler(object):
    """
    并行数据采样器
    """

    def __init__(self, User, usernum, itemnum, batch_size=64, maxlen=10, n_workers=1):
        """
        初始化采样器

        Args:
            User: 用户数据
            usernum: 用户数量
            itemnum: 物品数量
            batch_size: 批次大小
            maxlen: 序列最大长度
            n_workers: 工作进程数量
        """
        self.result_queue = Queue(maxsize=n_workers * 10)
        self.processors = []
        for i in range(n_workers):
            self.processors.append(
                Process(
                    target=sample_function,
                    args=(
                        User,
                        usernum,
                        itemnum,
                        batch_size,
                        maxlen,
                        self.result_queue,
                        np.random.randint(2e9),
                    ),
                )
            )
            self.processors[-1].daemon = True
            self.processors[-1].start()

    def next_batch(self):
        """获取下一批数据"""
        return self.result_queue.get()

    def close(self):
        """关闭所有工作进程"""
        for p in self.processors:
            p.terminate()
            p.join()


def data_partition(fname):
    """
    将数据集划分为训练集、验证集和测试集

    Args:
        fname: 数据文件名

    Returns:
        [user_train, user_valid, user_test, usernum, itemnum]: 训练集、验证集、测试集及用户数量和物品数量
    """
    # 获取当前文件的目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    # 构建数据文件的完整路径
    data_file = os.path.join(os.path.dirname(current_dir), "data", f"{fname}.txt")

    usernum = 0
    itemnum = 0
    User = defaultdict(list)
    user_train = {}
    user_valid = {}
    user_test = {}

    # 获取用户-物品交互数据
    with open(data_file, "r") as f:
        for line in f:
            u, i = line.rstrip().split(" ")
            u = int(u)
            i = int(i)
            usernum = max(u, usernum)
            itemnum = max(i, itemnum)
            User[u].append(i)

    # 划分数据集
    for user in User:
        nfeedback = len(User[user])
        if nfeedback < 3:
            user_train[user] = User[user]
            user_valid[user] = []
            user_test[user] = []
        else:
            user_train[user] = User[user][:-2]
            user_valid[user] = []
            user_valid[user].append(User[user][-2])
            user_test[user] = []
            user_test[user].append(User[user][-1])
    return [user_train, user_valid, user_test, usernum, itemnum]


def evaluate(model, dataset, args):
    """
    在测试集上评估模型性能

    Args:
        model: 待评估的模型
        dataset: 数据集
        args: 参数配置

    Returns:
        (NDCG, HT): NDCG和命中率指标
    """
    [train, valid, test, usernum, itemnum] = copy.deepcopy(dataset)

    NDCG = 0.0
    HT = 0.0
    valid_user = 0.0

    # 如果用户数量过多，随机采样10000个用户进行评估
    if usernum > 10000:
        users = random.sample(range(1, usernum + 1), 10000)
    else:
        users = range(1, usernum + 1)

    for u in users:
        if len(train[u]) < 1 or len(test[u]) < 1:
            continue

        seq = np.zeros([args.maxlen], dtype=np.int32)
        idx = args.maxlen - 1
        seq[idx] = valid[u][0]
        idx -= 1
        for i in reversed(train[u]):
            seq[idx] = i
            idx -= 1
            if idx == -1:
                break

        rated = set(train[u])
        rated.add(0)
        item_idx = [test[u][0]]
        # 随机采样100个负样本
        for _ in range(100):
            t = np.random.randint(1, itemnum + 1)
            while t in rated:
                t = np.random.randint(1, itemnum + 1)
            item_idx.append(t)

        predictions = -model.predict(*[np.array(l) for l in [[u], [seq], item_idx]])
        predictions = predictions[0]

        rank = predictions.argsort().argsort()[0].item()

        valid_user += 1

        if rank < 10:
            NDCG += 1 / np.log2(rank + 2)
            HT += 1
        if valid_user % 100 == 0:
            print(".", end="")
            sys.stdout.flush()

    return NDCG / valid_user, HT / valid_user


def evaluate_valid(model, dataset, args):
    """
    在验证集上评估模型性能

    Args:
        model: 待评估的模型
        dataset: 数据集
        args: 参数配置

    Returns:
        (NDCG, HT): NDCG和命中率指标
    """
    [train, valid, test, usernum, itemnum] = copy.deepcopy(dataset)

    NDCG = 0.0
    valid_user = 0.0
    HT = 0.0

    # 如果用户数量过多，随机采样10000个用户进行评估
    if usernum > 10000:
        users = random.sample(range(1, usernum + 1), 10000)
    else:
        users = range(1, usernum + 1)

    for u in users:
        if len(train[u]) < 1 or len(valid[u]) < 1:
            continue

        seq = np.zeros([args.maxlen], dtype=np.int32)
        idx = args.maxlen - 1
        for i in reversed(train[u]):
            seq[idx] = i
            idx -= 1
            if idx == -1:
                break

        rated = set(train[u])
        rated.add(0)
        item_idx = [valid[u][0]]
        # 随机采样100个负样本
        for _ in range(100):
            t = np.random.randint(1, itemnum + 1)
            while t in rated:
                t = np.random.randint(1, itemnum + 1)
            item_idx.append(t)

        predictions = -model.predict(*[np.array(l) for l in [[u], [seq], item_idx]])
        predictions = predictions[0]

        rank = predictions.argsort().argsort()[0].item()

        valid_user += 1

        if rank < 10:
            NDCG += 1 / np.log2(rank + 2)
            HT += 1
        if valid_user % 100 == 0:
            print(".", end="")
            sys.stdout.flush()

    return NDCG / valid_user, HT / valid_user
